import numpy as np
from common.utils.format_util import dup_name_handler
import statsmodels.api as sm


def run(df, params):
    col = params.get("col")
    all_cols = df.columns.tolist()

    data = df[col].values
    valid_index = np.where(np.isnan(data) == 0)[0]
    valid_data = data[valid_index]
    if "period" in params:
        period = int(params.get("period"))
    else:
        period = 4
    rd = sm.tsa.seasonal_decompose(valid_data, period=period)

    seasonal = '_'.join([col, "seasonal"])
    seasonal = dup_name_handler(seasonal, all_cols)
    trend = '_'.join([col, "trend"])
    trend = dup_name_handler(trend, all_cols)
    resid = '_'.join([col, "resid"])
    resid = dup_name_handler(resid, all_cols)

    index = all_cols.index(col) + 1
    all_cols.insert(index, resid)
    all_cols.insert(index, trend)
    all_cols.insert(index, seasonal)
    df = df.reindex(columns=all_cols)
    df[seasonal] = np.full(len(data), np.nan)
    df[trend] = np.full(len(data), np.nan)
    df[resid] = np.full(len(data), np.nan)
    df[seasonal][valid_index] = rd.seasonal
    df[trend][valid_index] = rd.trend
    df[resid][valid_index] = rd.resid

    return df
